from machine import Pin, reset
from time import sleep
import usocket as socket
import network
import gc
import _thread
import time

# 配置信息
WIFI_SSID = "LAPTOP-D3RPTUAA 6358"
WIFI_PASSWORD = "19234949992.ch"
TCP_PORT = 8080  # 接收ESP32-CAM视频流的端口
HTTP_PORT = 9000 # 提供Web访问的端口

# 固定IP配置 - 与ESP32-CAM中设置的S3_IP保持一致
FIXED_IP = '192.168.137.100'
GATEWAY = '192.168.137.1'     # 网关地址，通常是路由器IP
SUBNET = '255.255.255.0'      # 子网掩码
DNS = '8.8.8.8'               # DNS服务器

# 全局变量
tcp_server = None
client_sockets = []  # 存储HTTP客户端连接
frame_buffer = b""   # 存储最新一帧图像
frame_lock = _thread.allocate_lock()  # 用于线程安全的锁
frame_count = 0      # 接收到的帧数
last_receive_time = 0  # 最后接收时间

# WiFi连接类（带固定IP配置）
class WiFiConnection:
    def __init__(self, ssid, password, fixed_ip, gateway, subnet, dns):
        self.ssid = ssid
        self.password = password
        self.fixed_ip = fixed_ip
        self.gateway = gateway
        self.subnet = subnet
        self.dns = dns
        self.wlan = network.WLAN(network.STA_IF)
        
    def connect(self):
        # 1. 关闭 AP
        network.WLAN(network.AP_IF).active(False)

        # 2. 激活 STA
        self.wlan.active(True)
        sleep(0.5)          # 给硬件一点初始化时间

        # 3. 设置静态 IP（一定要在 connect 之前）
        ip_config = (self.fixed_ip, self.subnet, self.gateway, self.dns)
        self.wlan.ifconfig(ip_config)
        print("静态IP已设置为:", ip_config)

        # 4. 开始连接
        if not self.wlan.isconnected():
            print(f"开始连接 WiFi: {self.ssid}")
            self.wlan.connect(self.ssid, self.password)

            for t in range(30, 0, -1):
                if self.wlan.isconnected():
                    break
                print(f"连接中... {t}s")
                sleep(1)

        if self.wlan.isconnected():
            print("WiFi 连接成功，ifconfig =", self.wlan.ifconfig())
            return self.wlan.ifconfig()[0]
        else:
            print("WiFi 连接失败")
            return None

# TCP服务器 - 接收ESP32-CAM的视频流
def tcp_server_task():
    global tcp_server, frame_buffer, frame_count, last_receive_time
    
    # 创建TCP服务器
    tcp_server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    tcp_server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    tcp_server.bind(('0.0.0.0', TCP_PORT))
    tcp_server.listen(1)
    print(f"TCP服务器启动，监听端口 {TCP_PORT}")
    
    while True:
        try:
            # 接受连接
            client, addr = tcp_server.accept()
            print(f"ESP32-CAM已连接: {addr}")
            client.settimeout(5.0)
            
            buffer = b""
            
            while True:
                # 接收数据
                data = client.recv(1024)
                if not data:
                    break
                    
                buffer += data
                last_receive_time = time.time()
                
                # 检查是否接收到完整的一帧
                if b'--frame' in buffer:
                    # 分割帧
                    parts = buffer.split(b'--frame', 1)
                    if len(parts) == 2:
                        frame_data = parts[0]
                        buffer = b'--frame' + parts[1]
                        
                        # 提取JPEG裸数据（去掉HTTP头）
                        if b'Content-Length:' in frame_data:
                            try:
                                # 找到JPEG数据起始位置（两个\r\n之后）
                                jpeg_start = frame_data.find(b'\r\n\r\n') + 4
                                jpeg_data = frame_data[jpeg_start:]
                                if jpeg_data:
                                    with frame_lock:
                                        frame_buffer = jpeg_data
                                        frame_count += 1
                            except Exception as e:
                                print("解析JPEG失败:", e)
            
            client.close()
            print(f"ESP32-CAM断开连接: {addr}")
            
        except Exception as e:
            print(f"TCP服务器错误: {e}")
            try:
                client.close()
            except:
                pass

# HTTP服务器 - 提供Web页面和视频流
def http_server_task(ip):
    global client_sockets
    
    # HTML页面
    html = b"""<!DOCTYPE html>
<html>
<head>
    <title>ESP32-S3 视频流</title>
    <style>
        body {
            font-family: Arial, sans-serif;
            text-align: center;
            margin: 0;
            padding: 20px;
            background-color: #f0f0f0;
        }
        .container {
            max-width: 1000px;
            margin: 0 auto;
        }
        h1 {
            color: #333;
        }
        #videoContainer {
            background-color: #000;
            border-radius: 8px;
            padding: 10px;
            margin-top: 20px;
        }
        #status {
            margin-top: 20px;
            padding: 10px;
            background-color: #fff;
            border-radius: 4px;
        }
    </style>
</head>
<body>
    <div class="container">
        <h1>ESP32 摄像头视频流</h1>
        <div id="videoContainer">
            <img src="/stream" style="max-width: 100%;" />
        </div>
        <div id="status">等待视频流…</div>
    </div>
</body>
</html>
"""
    
    # 创建HTTP服务器
    http_server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    http_server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    http_server.bind(('0.0.0.0', HTTP_PORT))
    http_server.listen(5)
    print(f"HTTP 服务器启动，访问 http://{ip}:{HTTP_PORT} 查看视频")
    
    while True:
        try:
            client, addr = http_server.accept()
            client.settimeout(5.0)
            request = client.recv(1024).decode('utf-8', 'ignore')

            if 'GET / ' in request or 'GET /index.html' in request:
                # 发送 HTML
                header = (
                            b"HTTP/1.1 200 OK\r\n"
                            b"Content-Type: text/html; charset=utf-8\r\n"
                            b"Connection: close\r\n"
                            b"Content-Length: %d\r\n\r\n" % len(html)
                            )
                client.send(header + html)
                client.close()

            elif 'GET /stream' in request:
                # 升级为 MJPEG 流
                header = b"HTTP/1.1 200 OK\r\n" \
                         b"Content-Type: multipart/x-mixed-replace; boundary=frame\r\n" \
                         b"Connection: keep-alive\r\n" \
                         b"Cache-Control: no-cache\r\n\r\n"
                client.send(header)

                # 加入客户端列表
                with frame_lock:
                    client_sockets.append(client)

            else:
                client.send(b"HTTP/1.1 404 Not Found\r\n\r\n")
                client.close()

        except Exception as e:
            print("[HTTP] 异常:", e)
            try:
                client.close()
            except:
                pass
        gc.collect()

# 视频流分发任务 - 向所有HTTP客户端发送视频帧
def stream_distributor_task():
    global frame_buffer, client_sockets, frame_count
    
    while True:
        try:
            # 检查是否有新帧
            current_frame = None
            with frame_lock:
                if frame_buffer:
                    current_frame = frame_buffer
                    frame_buffer = b""  # 清空缓冲区，准备接收新帧
            
            # 如果有新帧，发送给所有客户端
            if current_frame:
                # 复制客户端列表以避免迭代中修改
                with frame_lock:
                    clients = list(client_sockets)
                
                for client in clients:
                    try:
                        # 发送MJPEG帧格式
                        client.send(b"--frame\r\n")
                        client.send(b"Content-Type: image/jpeg\r\n")
                        client.send(f"Content-Length: {len(current_frame)}\r\n\r\n".encode())
                        client.send(current_frame)
                        client.send(b"\r\n")
                    except Exception as e:
                        print(f"发送给客户端失败: {e}")
                        try:
                            with frame_lock:
                                client_sockets.remove(client)
                            client.close()
                        except:
                            pass
            
            # 检查视频流状态并通过串口打印
            current_time = time.time()
            if current_time - last_receive_time > 5:  # 5秒内没有收到帧
                print("警告: 未接收到视频流")
                sleep(10)
            else:
                # 每30帧打印一次状态
                if frame_count % 30 == 0 and frame_count > 0:
                    print(f"视频流状态: 正常 | 已接收 {frame_count} 帧")
            
            sleep(0.05)  # 控制循环频率
            
        except Exception as e:
            print(f"流分发错误: {e}")
            sleep(1)

# 主函数
def main():
    global last_status_time
    
    print("ESP32-S3 视频流接收器启动中...")
    gc.enable()
    
    # 连接WiFi（带固定IP配置）
    wifi = WiFiConnection(
        WIFI_SSID, 
        WIFI_PASSWORD,
        FIXED_IP,
        GATEWAY,
        SUBNET,
        DNS
    )
    ip = wifi.connect()
    if not ip:
        print("WiFi连接失败，重启设备...")
        sleep(5)
        reset()
    
    # 初始化时间变量
    global last_receive_time, last_status_time
    last_receive_time = time.time()
    last_status_time = time.time()
    
    # 启动TCP服务器线程（接收视频流）
    _thread.start_new_thread(tcp_server_task, ())
    
    # 启动HTTP服务器线程（提供Web访问）
    _thread.start_new_thread(http_server_task, (ip,))
    
    # 启动视频流分发线程
    stream_distributor_task()  # 在主线程运行

if __name__ == "__main__":
    main()

